#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/DepthFirstIterator.h"
#include "llvm/ADT/GraphTraits.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Analysis/DomTreeUpdater.h"
#include "llvm/Analysis/GlobalsModRef.h"
#include "llvm/Analysis/IteratedDominanceFrontier.h"
#include "llvm/Analysis/PostDominators.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/CFG.h"
#include "llvm/IR/DebugInfoMetadata.h"
#include "llvm/IR/DebugLoc.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/InstrTypes.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicInst.h"
#include "llvm/IR/PassManager.h"
#include "llvm/IR/Use.h"
#include "llvm/IR/Value.h"
#include "llvm/InitializePasses.h"
#include "llvm/Pass.h"
#include "llvm/ProfileData/InstrProf.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Scalar.h"
#include <cassert>
#include <cstddef>
#include <utility>
#include <vector>
#include <set>
#include <queue>
#include <unordered_map>

using namespace llvm;
/*
enum ColorType
{
    WHITE = 0,
    GRAY,
    BLACK
};
*/
namespace llvm
{
    ModulePass *createFDPass();
    void initializeFDPassPass(PassRegistry &);
} // namespace llvm

namespace
{
    struct Info
    {
        enum ColorType Color;
        int f;  //DFS 结束时间
        Info(ColorType color, int f): Color(color), f(f) {}
    };

    struct FDPass : public ModulePass
    {
        static char ID; // Pass identification, replacement for typeid

        std::unordered_map<llvm::Function *, std::vector<llvm::Function *>> func_call_Graph;
        std::unordered_map<llvm::Function *, struct Info *> DFSInfo;
        std::vector<llvm::Function *> functions;
        int Time_Count = 0;

        FDPass() : ModulePass(ID)
        {
            initializeFDPassPass(*PassRegistry::getPassRegistry());
        }

        void TopologicalSort(llvm::Module &M)
        {
            Time_Count = 0;
            for (Function &F : M.getFunctionList())
            {
                if (DFSInfo[&F]->Color == ColorType::WHITE)
                {
                    DFS(&F);
                }
            }
        }

        void DFS(Function *F)
        {
            DFSInfo[F]->Color = ColorType::GRAY;
            for (auto Called_F : func_call_Graph[F])
            {
                if (DFSInfo[Called_F]->Color == ColorType::WHITE)
                {
                    DFS(Called_F);
                }
                else if (DFSInfo[Called_F]->Color == ColorType::GRAY)
                {//存在一条后向边, 报告循环调用
                    llvm::outs() << "There exists a Circular Call between " << Called_F->getName() << "() and " << F->getName() << "() !\n";
                }
            }
            DFSInfo[F]->Color = ColorType::BLACK;
            DFSInfo[F]->f = ++Time_Count;
        }
/*
        static bool Compare_f(std::pair<llvm::Function *const, Info *> *Pair1, std::pair<llvm::Function *const, Info *> *Pair2)
        {
            return (Pair1->second->f > Pair2->second->f);
        }
*/
        bool runOnModule(llvm::Module &M) override
        {
            if (skipModule(M))
            {
                return false;
            }
            functions.clear();
            for (Function &F : M.getFunctionList())    //统计在每个函数中, 被调用的函数
            {
                functions.push_back(&F);
                func_call_Graph[&F].clear();
                DFSInfo[&F] = new Info(ColorType::WHITE, 0);

                for (auto I = inst_begin(&F), end = inst_end(&F); I != end; ++I)
                {
                    if (isa<CallInst>(*I))
                    {
                        auto Called_Func = cast<CallInst>(*I).getCalledFunction();
                        if (Called_Func != &F)  //不考虑递归
                        {
                            func_call_Graph[&F].push_back(Called_Func);
                        }
                    }
                }
            }

            TopologicalSort(M);
            //std::sort(DFSInfo.begin(), DFSInfo.end(), Compare_f);

            for (auto F = functions.begin(), end = functions.end(); F != end; ++F)
            {
                llvm::outs() << "Function " << (*F)->getName() << "() Called Function: ";
                if (func_call_Graph[*F].empty())
                {
                    llvm::outs() << "NIL";
                }
                for (auto Called_F : func_call_Graph[(*F)])
                {
                    llvm::outs() << Called_F->getName() << "()\t";
                }
                llvm::outs() << "\n";
            }

            return true;
        }

        void getAnalysisUsage(AnalysisUsage &AU) const override
        {
            AU.setPreservesCFG();
        }
    };
}

char FDPass::ID = 0;
INITIALIZE_PASS(FDPass, "FDPass", "Function Dependency", false, false)

ModulePass *llvm::createFDPass() { return new FDPass(); }